"""
--------------------------------------------------------------------------------
功能:读取陀螺仪原始数据, 并计算得到陀螺仪的滚转角, 俯仰角
日期:2025-06-03
作者:Unirobot
QQ交流群:811348489
寄语：代码中能遇见你，真好，有你，真好.
署名：有你同创智能机器人科技(北京)有限公司
--------------------------------------------------------------------------------
"""

import math
import json
import time
from collections import deque

value1 = 0.0  # 滚转角初始值
value2 = 0.0  # 俯仰角初始值
content_template = """
.----------------------------.
|                             |
|                             |
|        陀螺仪校准完成       |
|                             |
|                             |
.----------------------------.
滚转角初始值:{val1:<15}
俯仰角初始值:{val2:<15}
"""

class ProcessGyro:
    def __init__(self):
        self.alpha = 0.8 #互补滤波器系数
        self.window_size = 2  #移动平均窗口大小

        #初始化状态变量
        self.pitch_acc = 0.0 #加速度计计算的俯仰角
        self.roll_acc = 0.0  #加速度计计算的滚转角
        self.pitch = 0.0     #融合后的俯仰角
        self.roll = 0.0      #融合后的滚转角

        #时间相关变量
        self.last_time = time.time()
        self.current_time = time.time()
        self.gyro_x_history = deque(maxlen=self.window_size)
        self.gyro_y_history = deque(maxlen=self.window_size)
        self.gyro_z_history = deque(maxlen=self.window_size)

        # 校准数据
        self.calibration_samples = 50
        self.calibration_complete = False
        self.gyro_bias = {'x': 0, 'y': 0, 'z': 0}
        self.accel_bias = {'x': 0, 'y': 0, 'z': 0}
        self.calibration_data = {'gyro': [], 'accel': []}

        self.roll_add = []
        self.pitch_add = []
        self.average_roll = 0
        self.average_pitch = 0
        self.calculate_numbers = 0

        self.yaw = 0
        self.last_time = None
        self.sum_delta_time = 0.0

    """
    打印带颜色的字体
    """
    def color_text(self, text, color=None, bg_color=None, bold=False, underline=False):
        codes = []
        if color:
            color_codes = {
                "black": 30, "red": 31, "green": 32, "yellow": 33,
                "blue": 34, "magenta": 35, "cyan": 36, "white": 37,
            }
            codes.append(color_codes.get(color, 39))
        if bg_color:
            bg_codes = {
                "black": 40, "red": 41, "green": 42, "yellow": 43,
                "blue": 44, "magenta": 45, "cyan": 46, "white": 47,
            }
            codes.append(bg_codes.get(bg_color, 49))
        if bold:
            codes.append(1)
        if underline:
            codes.append(4)
        
        if codes:
            return f"\033[{';'.join(map(str, codes))}m{text}\033[0m"
        else:
            return text

    #################################################################
    #方法一：直接计算陀螺仪数据(目前该项目使用的是此方法)
    """
    陀螺仪标定: 计算出50组数据, 求和平均
    """
    def calibration_gyro(self, accel_data):
        global value1, value2
        self.calculate_numbers = self.calculate_numbers + 1
        if self.calculate_numbers <= self.calibration_samples:
            roll, pitch = self.calculate_angles(accel_data)
            self.roll_add.append(roll)
            self.pitch_add.append(pitch)

            dots = '.' * self.calculate_numbers
            colored_text = self.color_text(f"陀螺仪模块标定中{dots}", color="green", bold=True)
            print(colored_text, end='\r') #不换行打印数据
            # print(colored_text) #换行打印数据
            time.sleep(0.1)
        else:
            self.calibration_complete = True
            time.sleep(1)
            self.average_roll = sum(self.roll_add) / len(self.roll_add)
            self.average_pitch = sum(self.pitch_add) / len(self.pitch_add)
            # print("roll平均值:", self.average_roll, "pitch平均值:", self.average_pitch)
            content = content_template.format(
                val1=f"{self.average_roll:.2f}°", 
                val2=f"{self.average_pitch:.2f}°"
            )
            time.sleep(1)
            print(self.color_text(content, color="green"))
            time.sleep(2)

    """
    计算陀螺仪的数据
    """
    def calculate_angles(self, accel_data):
        """计算滚转角和俯仰角"""
        x, y, z = accel_data['x'], accel_data['y'], accel_data['z']
        roll = math.atan2(y, math.sqrt(x**2 + z**2)) * 180/math.pi
        pitch = math.atan2(-x, math.sqrt(y**2 + z**2)) * 180/math.pi
        roll = roll - self.average_roll
        pitch = pitch - self.average_pitch
        return roll, pitch

    def calculate_yaw(self, gyro_data, current_time=None):
        # 初始化
        if not hasattr(self, 'last_time'):
            self.last_time = current_time
            self.yaw = 0.0
            return 0.0
        
        # 计算时间间隔
        dt = current_time - self.last_time
        self.last_time = current_time
        
        # 更新偏航角
        self.yaw += gyro_data['z'] * dt
        
        # 漂移补偿
        if hasattr(self, 'yaw_drift_rate'):
            self.yaw -= self.yaw_drift_rate * dt * 100
        
        return self.yaw

    # """
    # 计算陀螺仪的数据(包含偏航角)
    # """
    # def calculate_angles_back(self, accel_data, gyro_data, current_time=None):
    #     """计算滚转角、俯仰角和偏航角"""
    #     # 计算滚转角和俯仰角
    #     x, y, z = accel_data['x'], accel_data['y'], accel_data['z']
    #     roll = math.atan2(y, math.sqrt(x**2 + z**2)) * 180/math.pi
    #     pitch = math.atan2(-x, math.sqrt(y**2 + z**2)) * 180/math.pi
    #     roll = roll - self.average_roll
    #     pitch = pitch - self.average_pitch
        
    #     # 计算偏航角(需要时间信息进行积分)
    #     if current_time is not None and self.last_time is not None:
    #         dt = current_time - self.last_time
    #         # print("dt:", dt)
    #         self.yaw += gyro_data['z'] * dt  # 对Z轴角速度积分
    #     self.last_time = current_time
    #     yaw_value = self.calculate_yaw(gyro_data, current_time)
        
    #     # return roll, pitch, self.yaw
    #     return roll, pitch, yaw_value

    """
    计算陀螺仪的数据(包含偏航角)
    """
    def calculate_angles_back(self, accel_data, gyro_data, current_time=None):
        """计算滚转角、俯仰角和偏航角"""
        # 计算滚转角和俯仰角
        x, y, z = accel_data['x'], accel_data['y'], accel_data['z']
        roll = math.atan2(y, math.sqrt(x**2 + z**2)) * 180/math.pi
        pitch = math.atan2(-x, math.sqrt(y**2 + z**2)) * 180/math.pi
        roll = roll - self.average_roll
        pitch = pitch - self.average_pitch
        
        # 计算偏航角(需要时间信息进行积分)
        if current_time is not None and self.last_time is not None:
            dt = current_time - self.last_time
            self.sum_delta_time += dt
            # print("dt:", dt)
            self.yaw += gyro_data['z'] * dt  # 对Z轴角速度积分
        self.last_time = current_time
        yaw_value = self.calculate_yaw(gyro_data, current_time)

        if self.sum_delta_time >= 0.6:
            self.sum_delta_time = 0
            self.yaw = 0
        
        # return roll, pitch, self.yaw
        # print("sum_dt:", self.sum_delta_time, yaw_value)
        return roll, pitch, yaw_value, self.sum_delta_time

def compensate_yaw(raw_yaw, elapsed_time):
    # 补偿参数通过线性回归获得
    drift_rate = -0.007  # 度/秒
    return raw_yaw - (drift_rate * elapsed_time)

    #################################################################
    #方法二：校准+计算陀螺仪数据
    """
    陀螺仪校准
    """
    def gyro_calibration(self, accel_data, gyro_data):
        # 如果没有完成校准，先进行校准
        if not self.calibration_complete:
            # 收集校准数据
            self.calibration_data['gyro'].append(gyro_data)
            self.calibration_data['accel'].append(accel_data)
            
            # 如果收集到足够样本，计算偏置
            if len(self.calibration_data['gyro']) >= self.calibration_samples:
                # 计算陀螺仪偏置
                gyro_x = [g['x'] for g in self.calibration_data['gyro']]
                gyro_y = [g['y'] for g in self.calibration_data['gyro']]
                gyro_z = [g['z'] for g in self.calibration_data['gyro']]
                
                self.gyro_bias['x'] = sum(gyro_x) / len(gyro_x)
                self.gyro_bias['y'] = sum(gyro_y) / len(gyro_y)
                self.gyro_bias['z'] = sum(gyro_z) / len(gyro_z)
                
                # 计算加速度计偏置
                accel_x = [a['x'] for a in self.calibration_data['accel']]
                accel_y = [a['y'] for a in self.calibration_data['accel']]
                accel_z = [a['z'] for a in self.calibration_data['accel']]
                
                self.accel_bias['x'] = sum(accel_x) / len(accel_x)
                self.accel_bias['y'] = sum(accel_y) / len(accel_y)
                # 理想情况下z轴应为1g（重力加速度）
                self.accel_bias['z'] = (sum(accel_z) / len(accel_z)) - 1.0
                
                print(f"[CALIBRATION] 校准完成. 陀螺仪偏置: {self.gyro_bias}")
                print(f"[CALIBRATION] 加速度计偏置: {self.accel_bias}")
                self.calibration_complete = True
            else:
                # 显示校准进度
                progress = len(self.calibration_data['gyro'])
                print(f"[CALIBRATION] 校准中... ({progress}/{self.calibration_samples})")
                return 0, 0

    """
    实时计算陀螺仪数据
    """
    def update_angles(self, accel_data, gyro_data):
        """
        使用新数据更新角度计算
        :param accel_data: 加速度计数据字典 {'x': x, 'y': y, 'z': z}
        :param gyro_data: 陀螺仪数据字典 {'x': x, 'y': y, 'z': z}
        :return: 更新后的角度 (roll, pitch)
        """
        # 应用校准
        ax = accel_data['x'] - self.accel_bias['x']
        ay = accel_data['y'] - self.accel_bias['y']
        az = accel_data['z'] - self.accel_bias['z']
        
        gx = gyro_data['x'] - self.gyro_bias['x']
        gy = gyro_data['y'] - self.gyro_bias['y']
        gz = gyro_data['z'] - self.gyro_bias['z']
        
        # 存储陀螺仪数据用于平滑
        self.gyro_x_history.append(gx)
        self.gyro_y_history.append(gy)
        self.gyro_z_history.append(gz)
        
        # 使用移动平均平滑陀螺仪数据
        gx = sum(self.gyro_x_history) / len(self.gyro_x_history)
        gy = sum(self.gyro_y_history) / len(self.gyro_y_history)
        gz = sum(self.gyro_z_history) / len(self.gyro_z_history)
        
        # 计算时间差
        self.current_time = time.time()
        dt = self.current_time - self.last_time
        self.last_time = self.current_time
        
        # 1. 使用加速度计计算角度
        # 计算滚转角 (绕X轴旋转)
        self.roll_acc = math.atan2(ay, math.sqrt(ax**2 + az**2)) * (180 / math.pi)
        
        # 计算俯仰角 (绕Y轴旋转)
        self.pitch_acc = math.atan2(-ax, math.sqrt(ay**2 + az**2)) * (180 / math.pi)
        
        # 2. 使用陀螺仪计算角度变化
        # 转换为弧度/秒
        gx_rad = gx * (math.pi / 180)
        gy_rad = gy * (math.pi / 180)
        
        # 使用互补滤波器融合数据
        self.pitch = self.alpha * (self.pitch + gy_rad * dt) + (1 - self.alpha) * self.pitch_acc
        self.roll = self.alpha * (self.roll + gx_rad * dt) + (1 - self.alpha) * self.roll_acc
        return self.roll, self.pitch